package com.example.mybatisplus.threadPool;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.concurrent.ListenableFuture;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * @description: 重写线程池，传递线程变量
 * @author: dengPing.xu
 * @date: 2022-03-19
 */
public class InheritableThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {
    @Override
    public void execute(final Runnable task) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        super.execute(() -> execute(task, userId));
    }

    @Override
    public void execute(final Runnable task, long startTimeout) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        super.execute(() -> execute(task, userId), startTimeout);
    }

    @Override
    public Future<?> submit(final Runnable task) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        return super.submit(() -> execute(task, userId));
    }

    @Override
    public <T> Future<T> submit(final Callable<T> task) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        return super.submit(() -> execute(task, userId));
    }

    @Override
    public ListenableFuture<?> submitListenable(final Runnable task) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        return super.submitListenable(() -> execute(task, userId));
    }

    @Override
    public <T> ListenableFuture<T> submitListenable(final Callable<T> task) {
        final Integer userId = SpdUserIdHolder.getCurrentUserId();
        return super.submitListenable(() -> execute(task, userId));
    }

    private void execute(Runnable task, Integer userId) {
        SpdUserIdHolder.set(userId);
        task.run();
        SpdUserIdHolder.clear();
    }

    private <T> T execute(Callable<T> task, Integer userId) throws Exception {
        SpdUserIdHolder.set(userId);
        T value = task.call();
        SpdUserIdHolder.clear();
        return value;
    }
}
